// Copyright (c) 2023 Presto Labs Pte. Ltd.
// Author: jhkim

#include "cc/appcoin2/research/util/option_util.h"
#include <iostream>
#include <cmath>
#include <vector>
#include <algorithm>
#include <tuple>
#include <limits>
#include <iomanip>

#include "util/lbfgs/lbfgs.h"

#include <Eigen/Dense>
#include <unsupported/Eigen/LevenbergMarquardt>


// Include the L-BFGS header here
// You can find the header file and its source file at https://github.com/imneme/pcg-c-basic
// Download the files lbfgs.h and lbfgs.c and include them in your project

// Rest of the code


double round_to_nth(double value, int decimals) {
  double factor = std::pow(10, decimals);
  return std::round(value * factor) / factor;
}

// SABR model functions
double SABR(double F, double K, double T, double alpha, double beta, double rho, double nu) {
  double sabrsigma;

  if (fabs(F - K) < 0.00001) {
    double numer1 = pow((1 - beta), 2) / 24 * pow(alpha, 2) / pow(F, (2 - 2 * beta));
    double numer2 = 0.25 * rho * beta * nu * alpha / pow(F, (1 - beta));
    double numer3 = (2 - 3 * pow(rho, 2)) / 24 * pow(nu, 2);
    sabrsigma = alpha * (1 + (numer1 + numer2 + numer3) * T) / pow(F, (1 - beta));
  }
  else {
    double logFK = std::log(F / K);
    double z = (nu / alpha) * pow((F * K), ((1 - beta) / 2)) * logFK;
    double zhi = std::log((std::sqrt(1 - 2 * rho * z + pow(z, 2)) + z - rho) / (1 - rho));
    double numer1 = pow((1 - beta), 2) / 24 * (pow(alpha, 2) / pow((F * K), (1 - beta)));
    double numer2 = 0.25 * rho * beta * nu * alpha / pow((F * K), ((1 - beta) / 2));
    double numer3 = (2 - 3 * pow(rho, 2)) / 24 * pow(nu, 2);
    double numer = alpha * (1 + (numer1 + numer2 + numer3) * T) * z;
    double denom1 = pow((1 - beta), 2) / 24 * pow(logFK, 2);
    double denom2 = pow((1 - beta), 4) / 1920 * pow(logFK, 4);
    double denom = pow((F * K), ((1 - beta) / 2)) * (1 + denom1 + denom2) * zhi;
    sabrsigma = numer / denom;
  }

  return sabrsigma;
}

// Objective function for calibration
static lbfgsfloatval_t SABRcalibration(void *instance, const lbfgsfloatval_t *x, lbfgsfloatval_t *g, const int n, const lbfgsfloatval_t step) {
  auto *data = reinterpret_cast<std::tuple<std::vector<double>, std::vector<double>, double, double, double>*>(instance);

  auto& strikes = std::get<0>(*data);
  auto& vols = std::get<1>(*data);
  double beta = std::get<2>(*data);
  double F = std::get<3>(*data);
  double T = std::get<4>(*data);

  double err = 0.0;
  for (size_t i = 0; i < vols.size(); ++i) {
    double diff = vols[i] - SABR(F, strikes[i], T, x[0], beta, x[1], x[2]);
    err += diff * diff;
    LOG(INFO) << strikes[i] << " " << vols[i] << " " << SABR(F, strikes[i], T, x[0], beta, x[1], x[2]);
  }

  LOG(INFO) << x[0];
  LOG(INFO) << beta << " " << F << " " << T;
  LOG(INFO) << err;
  return err;
}

using namespace Eigen;

struct IvFunctor : DenseFunctor<double> {
  IvFunctor(
      double price_option,
      bool is_call,
      double S,
      double K,
      double r,
      double T)
      : DenseFunctor<double>(1, 1),
        price_option(price_option),
        is_call(is_call),
        S(S),
        K(K),
        r(r),
        T(T) {}
  double price_option;
  bool is_call;
  double S;
  double K;
  double r;
  double T;

  int operator()(const Eigen::VectorXd &x, Eigen::VectorXd &fvec) const {
    double diff = is_call ? call_price(S, K, r, T, x[0]) : put_price(S, K, r, T, x[0])
        - price_option;
    fvec[0] = diff;
    return 0;
  }
  int df(const Eigen::VectorXd &x, Eigen::MatrixXd &fjac) const {
    double h = 1e-8; // Small value for finite differences
    size_t i = 0;
    size_t j = 0;
    double prc_plus_h =
        is_call ? call_price(S, K, r, T, x[0] + h) : put_price(S, K, r, T, x[0] + h);
    double prc_minus_h =
        is_call ? call_price(S, K, r, T, x[0] - h) : put_price(S, K, r, T, x[0] - h);
    fjac(j, i) = (prc_plus_h - prc_minus_h) / (2.0 * h);
    return 0;
  }
};

double calibrate_iv(
    double price_option,
    bool is_call,
    double S,
    double K,
    double r,
    double T) {
  Eigen::VectorXd x(1); // Initial guess
  x << 0.5;
  IvFunctor functor{price_option, is_call, S, K, r, T};
  Eigen::LevenbergMarquardt<IvFunctor> lm(functor);
  lm.minimize(x);
  return x[0];
}


struct SabrFunctor : DenseFunctor<double> {
  SabrFunctor(
      std::vector<double> strikes,
      std::vector<double> midvol,
      double F,
      double T,
      double beta)
      : DenseFunctor<double>(3, strikes.size()),
        strikes(strikes),
        midvol(midvol),
        F(F),
        T(T),
        beta(beta) {}
  std::vector<double> strikes;
  std::vector<double> midvol;
  const double F;
  const double T;
  const double beta;

  int operator()(const Eigen::VectorXd &x, Eigen::VectorXd &fvec) const {
    for (size_t i = 0; i < midvol.size(); ++i) {
      double diff = SABR(F, strikes[i], T, x[0], beta, x[1], x[2]) - midvol[i];
      fvec[i] = diff;
      // LOG(INFO) << strikes[i] << " " << midvol[i] << " " << SABR(F, strikes[i], T, x[0], beta, x[1], x[2]);
    }
    return 0;
  }
  int df(const Eigen::VectorXd &x, Eigen::MatrixXd &fjac) const {
    double h = 1e-8; // Small value for finite differences
    Eigen::VectorXd x_plus_h = x;
    Eigen::VectorXd x_minus_h = x;
    for (size_t i = 0; i < x.size(); ++i) {
      x_plus_h[i] += h;
      x_minus_h[i] -= h;

      for (size_t j = 0; j < strikes.size(); ++j) {
        double sabr_plus_h = SABR(F, strikes[j], T, x_plus_h[0], beta, x_plus_h[1], x_plus_h[2]);
        double sabr_minus_h = SABR(F, strikes[j], T, x_minus_h[0], beta, x_minus_h[1], x_minus_h[2]);
        fjac(j, i) = (sabr_plus_h - sabr_minus_h) / (2.0 * h);
      }

      x_plus_h[i] -= h;
      x_minus_h[i] += h;
    }

    return 0;
  }
};


// Function to calibrate SABR model parameters
std::tuple<double, double, double> calibrate_sabr(double f, std::vector<double> &strikes, double t, std::vector<double> &midvol) {
  Eigen::VectorXd x(3); // Initial guess
  x << 0.4, -0.1, 1.2;

  SabrFunctor functor{strikes, midvol, f, t, 1};
  Eigen::LevenbergMarquardt<SabrFunctor> lm(functor);
  lm.minimize(x);

  return std::make_tuple(x[0], x[1], x[2]);
}

auto unused(double f, std::vector<double> &strikes, double t, std::vector<double> &midvol) {
  std::vector<double> initial_guess = {0.4, -0.1, 1.2};

  lbfgs_parameter_t param;
  lbfgs_parameter_init(&param);
  param.epsilon = 1e-6;

  int n = 3; // number of parameters to optimize
  lbfgsfloatval_t fx;
  lbfgsfloatval_t *x = lbfgs_malloc(n);
  for (int i = 0; i < n; i++) {
    x[i] = initial_guess[i];
  }

  auto tpl = std::make_tuple(strikes, midvol, 1.0, f, t);
  int ret = lbfgs(n, x, &fx, SABRcalibration, nullptr, &tpl, &param);

  printf("L-BFGS optimization terminated with status code = %d\n", ret);

  double alpha = round_to_nth(x[0], 4);
  double rho = round_to_nth(x[1], 4);
  double nu = round_to_nth(x[2], 4);

  lbfgs_free(x);

  return std::make_tuple(alpha, rho, nu);
}

inline std::vector<std::string> Split(std::string_view sv, std::string_view delims) {
  std::vector<std::string> out;
  while (true) {
    if (sv.empty()) {
      out.emplace_back(sv);
      break;
    }
    int idx = sv.find_first_of(delims);
    if (idx == std::string_view::npos) {
      out.emplace_back(sv);
      break;
    } else {
      out.emplace_back(sv.substr(0, idx));
      sv = sv.substr(idx + 1);
    }
  }
  return out;
}

#include <optional>
#include <chrono>

std::tm int_to_tm(int date_int) {
  int year = date_int / 10000;
  int month = (date_int % 10000) / 100;
  int day = date_int % 100;
  std::tm date_tm = {};
  date_tm.tm_year = year - 1900;
  date_tm.tm_mon = month - 1;
  date_tm.tm_mday = day;
  return date_tm;
}

double duration_in_year(int64_t timestamp_ns, int64_t date_int) {
  // Convert the timestamp to a time_point
  std::chrono::nanoseconds ns(timestamp_ns);
  std::chrono::system_clock::time_point tp(std::chrono::duration_cast<std::chrono::system_clock::duration>(ns));
  // Convert the integer date to a tm struct
  std::tm date_tm = int_to_tm(date_int);
  // Convert the tm struct to a time_point
  std::chrono::system_clock::time_point date_tp = std::chrono::system_clock::from_time_t(std::mktime(&date_tm));
  // Calculate the duration in nanoseconds between the timestamp and the date
  std::chrono::nanoseconds duration = date_tp - tp;
  return duration.count() / 1e9 / 3600. / 24. / 365.;
}

struct OptionCP {
  int64_t exp_ymd;
  std::string strike;
  double index = 0;
  double expiry = std::numeric_limits<double>::quiet_NaN();
  double price_call = std::numeric_limits<double>::quiet_NaN();
  double price_put = std::numeric_limits<double>::quiet_NaN();
};

struct OptionSet {
  int64_t exp_ymd;
  std::map<std::string, OptionCP> options_sameexp = {};
  OptionCP& get(const std::string& strike) {
    if (options_sameexp.count(strike) == 0) {
      options_sameexp[strike] = OptionCP{exp_ymd, strike};
    }
    return options_sameexp.at(strike);
  }
  void PrintString() {
    std::vector<double> strikes;
    std::vector<double> vols;
    double r = 0.02;
    double F = 0;
    double expiry = 0;
    for (auto& optpair : options_sameexp) {
      auto& opt = optpair.second;
      expiry = opt.expiry;
      if (!std::isnan(opt.price_call) || !std::isnan(opt.price_put)) {
        // std::cout << exp_ymd << " " << opt.strike << " : " << opt.index << " C" << opt.price_call << " P" << opt.price_put << " | ";
      }
      double price_option = 0;
      bool is_call = false;
      if (opt.index > std::stod(opt.strike)) {
        price_option = opt.price_put;
      } else {
        price_option = opt.price_call;
        is_call = true;
      }
      double strike = std::stod(opt.strike);
      if (!std::isnan(price_option)) {
        double vol = calibrate_iv(
            price_option * opt.index,
            is_call,
            opt.index,
            strike,
            r,
            opt.expiry);
        strikes.push_back(strike);
        vols.push_back(vol);
        LOG(INFO) << opt.expiry << " " << strike << ", " << vol;
      }
      F = opt.index * std::exp(r * opt.expiry);
    }

    double alpha, rho, nu;
    std::tie(alpha, rho, nu) = calibrate_sabr(F, strikes, expiry, vols);
    LOG(INFO) << "SABR parameters:" << std::endl;
    LOG(INFO) << "Alpha: " << alpha << ", Rho: " << rho << ", Nu: " << nu << std::endl;
  }
};

struct OptionsSurface {
  OptionCP& GetMutableOptionCP(std::string instrument_name) {
    std::vector<std::string> splitted = Split(instrument_name, ".");
    return GetMutableOptionCP(std::stoll(splitted[1]), splitted[3]);
  }
  OptionCP& GetMutableOptionCP(
      int64_t exp_ymd,
      const std::string& strike) {
    return get(exp_ymd).get(strike);
  }
  OptionSet& get(int64_t exp_ymd) {
    if (options_all.count(exp_ymd) == 0) {
      options_all[exp_ymd] = OptionSet{exp_ymd};
    }
    return options_all.at(exp_ymd);
  }
  void Update(
      std::string symbol,
      int64_t timestamp,
      double index,
      double midp) {
    auto& ocp = GetMutableOptionCP(symbol);
    if (symbol.find(".C.") <= symbol.size()) {
      ocp.price_call = midp;
    } else if (symbol.find(".P.") <= symbol.size()) {
      ocp.price_put = midp;
    }
    ocp.index = index;
    ocp.expiry = duration_in_year(timestamp, ocp.exp_ymd);
    // LOG(INFO) << symbol << " " << timestamp << " " << index << " " << midp << " " << ocp.expiry;
  }
  void PrintString() {
    for (auto& opts : options_all) {
      opts.second.PrintString();
    }
  }
  std::map<int64_t, OptionSet> options_all;
};

int example() {
  double F = 30079.71; // Forward price, this should be obtained from market data
  std::vector<double> strikes = {31000.0, 36000.0, 40000.0, 35000.0, 32000.0, 38000.0, 34000.0, 29000.0, 24000.0, 26000.0, 23000.0, 27000.0, 21000.0, 28000.0, 25000.0, 30000.0, 22000.0};
  double t = 0.041134589580162355; // Time to maturity in years, this should be obtained from market data
  std::vector<double> midvol = {0.53, 0.6, 0.69, 0.58, 0.53, 0.65, 0.57, 0.49, 0.66, 0.57, 0.73, 0.53, 0.82, 0.51, 0.61, 0.49, 0.78};

  double alpha, rho, nu;
  std::tie(alpha, rho, nu) = calibrate_sabr(F, strikes, t, midvol);

  LOG(INFO) << "SABR parameters:" << std::endl;
  LOG(INFO) << "Alpha: " << alpha << ", Rho: " << rho << ", Nu: " << nu << std::endl;

  double sabriv = SABR(F, 29500, t, alpha, 1, rho, nu);
  LOG(INFO) << "SABR implied volatility for strike 29500: " << sabriv << std::endl;

  return 0;
}

#include <experimental/filesystem>
#include <highfive/H5File.hpp>
#include <highfive/H5DataSet.hpp>
#include <highfive/H5DataSpace.hpp>


std::vector<std::vector<double>> read_matrix(
    HighFive::File file,
    const std::string& column) {
  std::vector<std::vector<double>> ret;
  HighFive::DataSet dataset = file.getDataSet(column);
  HighFive::DataSpace dataspace = dataset.getSpace();
  // Get the dimensions of the dataset
  std::vector<size_t> dims = dataspace.getDimensions();
  // Allocate memory for the data
  std::vector<double> data(dims[0] * dims[1]);
  // Read the dataset data into memory
  dataset.read(data.data());
  ret.resize(dims[0]);
  for (size_t i = 0; i < dims[0]; ++i) {
    ret[i].resize(dims[1]);
    for (size_t j = 0; j < dims[1]; ++j) {
      ret[i][j] = data[i * dims[1] + j];
    }
  }
  return ret;
}

template<class T>
std::vector<T> read_row(HighFive::File file, const std::string& column) {
  HighFive::DataSet dataset = file.getDataSet(column);
  std::vector<T> data(dataset.getSpace().getDimensions()[0]);
  dataset.read(data);
  return data;
}

int main(int argc, char* argv[]) {
  InitializeTracer();
  // example();
  cxxopts::Options opt("option_calculator", "options calculator");
  opt.add_options()("trading_date", "", cxxopts::value<std::string>());
  opt.add_options()("ccy", "BTC", cxxopts::value<std::string>());
  auto res = opt.parse(argc, argv);
  auto yyyymmdd = res["trading_date"].as<std::string>();
  auto ccy = res["ccy"].as<std::string>();

  std::experimental::filesystem::path root_dir =
      "/remote/iosg/data-2/buckets/feed.derived.interval_h5/coin/main/PT1M";
  auto path = root_dir.append(
      "Options.Deribit.v2").append(yyyymmdd).append("Options.Deribit.v2--ohlc.h5");

  HighFive::File file(path, HighFive::File::ReadOnly);

  std::vector<std::string> univdata = read_row<std::string>(file, "/universe");
  std::vector<int64_t> times = read_row<int64_t>(file, "/timestamp");
  auto indexs = read_matrix(file, "/CLOSE_INDEX");
  auto midps = read_matrix(file, "/CLOSE_MID");

  OptionsSurface osurf;

  for (int i = 0; i < univdata.size(); ++i) {
    auto symbol = univdata[i];
    auto timestamp = times.back();
    auto index = indexs.back()[i];
    auto midp = midps.back()[i];
    if (symbol.find(ccy) != 0) continue;
    osurf.Update(symbol, timestamp, index, midp);
  }

  osurf.PrintString();

  return 0;
}